import datetime
import traceback

from bson import ObjectId
from flask import current_app as app

from app.api.v2.cloud import system_add_black_list
from app.celery.handle_result.asset import AssetHandler
from app.celery.handle_result.available import AvailableHandler
from app.celery.handle_result.securityEvent import SecurityEventHandler
from app.celery.handle_result.ssl import SslHandler
from app.celery.handle_result.utils import update_task_alert_settings, save_spider_urls
from app.celery.handle_result.vul import VulHandler
from app.celery.handle_result.hostVul import vulHostHandler
from app.celery.handle_result.ipv6 import IPv6Handler
from app.celery.handle_result.changeCheck import ChangCheckHandler, update_baseline
from app.db import db, redis
from app.db.models.alert_groups import AlertGroup
from app.db.models.assets import Asset
from app.db.models.balances import Balances
from app.db.models.jobs import Job
from app.db.models.tasks import Task, Result, IgnoreWarning
from app.db.models.asset_tasks import AssetTasks
from app.db.models.api_token import OutPutConfig
from app.libs.alertsender import (
    VulAlertSender, AssetAlertSender,
    HttpAlertSender, PingAlertSender,
    SslAlertSender, ChangeCheckAlertSender,
    SecurityEventAlertSender, HostVulAlertSender
)
from app.libs.enums import TaskType, AVAILABLE_TASK_GROUP_TUPLE, WS_TASK_GROUP_TUPLE, AssetTaskGroup
from app.libs.utility import list_over_insert, list_protect_get
from app.api.v2.job_util import update_last_host_security_status, update_last_web_security_status
from app.config.settings import OUTPUT_QUEUE


class TaskResultHandler:
    """按任务类型对监测任务结果进行处理并告警"""

    def __init__(self, _db=None, retry=False):
        self.handler = {
            "http": AvailableHandler,
            "ping": AvailableHandler,
            "vul": VulHandler,
            "ssl": SslHandler,
            "securityEvent": SecurityEventHandler,
            "asset": AssetHandler,
            "hostVul": vulHostHandler,
            "ipv6": IPv6Handler,
            TaskType.change_check.value: ChangCheckHandler
        }
        self.db = _db if _db else db.get_db(alias="default")
        self.retry = retry

    def handle(self, task: Task.result):
        task_type = task.taskType
        app.logger.info(f"开始处理任务： Id: {task.id}, 任务类型：{task.taskType}, 资产地址：{task.target}")
        this_handler = self.handler.get(task_type)
        if this_handler:
            try:
                result = this_handler(task=task).handle()
                task.update(progress=100)
            except Exception as e:
                app.logger.error(f'[02]监测结果处理：task {task.id}异常:{e}')
                traceback.print_exc()
                return
            else:
                if not result:
                    return
                if task.status == "failed":
                    return self.handle_task_failed(task)
                else:
                    return self.handle_task_completed(task=task, result=result)

    @staticmethod
    def send_api_message(task):
        """推送任务信息"""
        out_config = OutPutConfig.objects.find({"uid": task.uid, "is_active": True}).first()
        if out_config and task.taskType in out_config.push_task_types:
            redis.rpush(OUTPUT_QUEUE, str(task.id))

    def handle_task_failed(self, task: Task):
        """处理失败监测任务"""
        app.logger.info(f"处理status==failed任务: {task.id}")
        try:
            task.status = "failed"
            task.result.warnings = []
            task.save()
            self.send_api_message(task)
        except Exception as e:
            app.logger.error(f'[01]更新结果失败：task {task.id}:{e}')
            return

    def handle_task_completed(self, task: Task, result: dict):
        """处理已完成的监测任务"""
        task_type = task.taskType
        task.isSpecialTask = task.triggerType == 'special'
        result = Result(**result)
        task.result = result
        task.isLatestTask = True
        # 将除当前任务的其他相同job的任务状态更新
        Task._get_collection().update_many(
            {
                "taskType": task.taskType,
                "isLatestTask": True,
                "_id": {"$ne": task.id},
                "jobId": task.jobId,
                "triggerType": {"$ne": 'special'}
            },
            {
                "$set": {"isLatestTask": False}
            }
        )
        ref = AssetTasks.objects.find_one({f"{task.taskType}Settings.taskId": task.taskId})
        job = Job.objects.find_one({"_id": task.jobId})
        asset = Asset.objects.find_one({"jobId": task.jobId})
        if not ref:
            app.logger.warning(f"处理任务完成逻辑：未查询到关联的Ref记录: task: {task.id}, refId: {task.refId}，删除该任务")
            task.delete()
            return

        task_idx = task.taskIdx
        if task.taskType == "http":
            task_idx = -1
            for index, item in enumerate(job.httpSettings):
                if item.target == task.target:
                    task_idx = index
                    break
        _warnings = result.warnings
        warnings = []
        ignore_warning_map = {w.detail.id: w.ignoredTime for w in IgnoreWarning.objects.filter(
            jobId=task.jobId, taskType=task.taskType)}
        for warning in _warnings:
            if not ignore_warning_map.get(warning.id):
                warning._id = ObjectId()
                warnings.append(warning)
        safe_result = []
        for warning in warnings:
            if job.ignoredWarnIdMap.get(warning.id):
                warning.ignoredTime = job.ignoredWarnIdMap.get(warning.id)
            this_is_safe = bool(warning.isResolved or warning.ignoredTime)
            safe_result.append(this_is_safe)
        task.result.warnings = warnings
        if len(warnings) and all(safe_result):
            task.securityStatus = "safe"

        if result.targetStatus and result.targetStatus.status != "good" and len(warnings) == 0:
            task.securityStatus = "unknown"
        else:
            task.securityStatus = "warning" if len(warnings) > 0 else "safe"
        for key, value in self.warning_counts_update(warnings).items():
            setattr(task.result, key, value)

        sec_status_field = f"{task.taskType}SecurityStatus"
        start_status_field = f"{task.taskType}IsStarted"
        result_id_field = f"{task.taskType}ResultId"
        abnormal_field = f"{task.taskType}Abnormal"
        prev_result_id_field = f"{task.taskType}PrevResultId"

        if task_type in AVAILABLE_TASK_GROUP_TUPLE:
            list_over_insert(job[sec_status_field] or [], task_idx, task.securityStatus, 'unknown')
            list_over_insert(job[start_status_field] or [], task.taskIdx, False, False)
        if task_type in AVAILABLE_TASK_GROUP_TUPLE + WS_TASK_GROUP_TUPLE + (TaskType.hostVul.value, ):
            ref[sec_status_field] = task.securityStatus

        if (task_type in WS_TASK_GROUP_TUPLE) or (task_type in AVAILABLE_TASK_GROUP_TUPLE):
            prev_result_id = ref[result_id_field]
            ref[result_id_field] = task.id
            if prev_result_id:
                setattr(ref, prev_result_id_field, prev_result_id)
            if task_type in WS_TASK_GROUP_TUPLE and prev_result_id:
                setattr(job, prev_result_id_field, prev_result_id)
            setattr(ref, abnormal_field, bool(task.securityStatus == "warning"))
            if task_type in WS_TASK_GROUP_TUPLE:
                setattr(job, abnormal_field, bool(task.securityStatus == "warning"))

        # 保留job数据结构
        if task_type in AVAILABLE_TASK_GROUP_TUPLE:

            list_over_insert(job[result_id_field] or [], task_idx, task.id, None)
            list_over_insert(job[abnormal_field] or [], task_idx, bool(task.securityStatus == "warning"), False)

            prev_result_id = list_protect_get(job[result_id_field], task_idx) if job[result_id_field] else None
            if prev_result_id:
                list_over_insert(job[prev_result_id_field], task_idx, prev_result_id, None)

        if (task.triggerType in ["schedule", "scheduleIncrement", "manual"] or
                task.taskSettings.enableSitePortraitTriggerMonitor):
            job, task = update_task_alert_settings(job, task)

        task_alert_enable = task.taskSettings.alert.get("enable")
        if task_alert_enable is not None and task.triggerType in ["special"]:
            task.alertSettings.enable[task_type] = task_alert_enable

        # 保存task job/asset 更新数据到数据库
        task.status = "completed"
        task.taskIdx = task_idx
        # 主机安全检测漏洞结果不用存在tasks表
        if task_type == "hostVul":
            task.result.warnings = []
        task.save(validate=False)
        if not self.retry:
            ref.save()
            job.save()

        # 如果是asset任务， 还需更新assets表数据
        if task_type == TaskType.asset.value and asset:
            asset["securityStatus"] = task.securityStatus
            asset["isStarted"] = False
            prev_result_id = asset.resultId
            asset.resultId = task.id
            if prev_result_id:
                asset.prevResultId = prev_result_id
            asset.save()

        # 更新job安全状态
        job = Job.objects.filter(pk=task.jobId).first()
        if task.taskType in WS_TASK_GROUP_TUPLE:
            update_last_web_security_status(job)
        elif task.taskType == TaskType.hostVul.value:
            update_last_host_security_status(job, task)

        try:
            save_spider_urls(task.id)
        except Exception as e:
            app.logger.info(
                f"[01]更新任务URL记录失败：job: {job.note}{job.id}, task: {task.target}-{task.taskType}-{task.id}, {e}")

        app.logger.info(f"[01]更新任务结果成功：job/asset: {job.note}{ref.id}, task: {task.target}-{task.taskType}-{task.id}")
        if isinstance(ref, AssetTasks):
            try:
                if ref.taskGroupType == AssetTaskGroup.web_security.value:
                    ref._get_web_security_last_result(is_save=True)
                elif ref.taskGroupType == AssetTaskGroup.host.value:
                    ref._save_host_last_result()
                ref.update_status(task_type=task.taskType, status=task.status)
            except:
                ref.lastResult = {}
                ref.save()

        app.logger.info(f"[01]更新任务结果成功：job/asset: {job.note}{ref.id}, task: {task.target}-{task.taskType}-{task.id}")

        if (task.taskType == TaskType.change_check.value and
                task.result.addition.get("statistics", {}).get("is_update_baseline")):
            try:
                update_baseline(task.triggerType, task.taskId, task.taskSessionId)
            except Exception as e:
                app.logger.error(f'[01]更新结果篡改监测样本失败：task {task.id}:{e}')

        # 告警发送
        if not self.retry and task.alertSettings.alertGroupIds:
            alert_group = AlertGroup.objects.filter(pk=task.alertSettings.alertGroupIds[0]).first()
            if alert_group and alert_group.enable:
                alert_sender_type_map = {
                    "vul": VulAlertSender,
                    "ssl": SslAlertSender,
                    "asset": AssetAlertSender,
                    "http": HttpAlertSender,
                    "ping": PingAlertSender,
                    "securityEvent": SecurityEventAlertSender,
                    "hostVul": HostVulAlertSender,
                    TaskType.change_check.value: ChangeCheckAlertSender
                }
                alert_sender_class = alert_sender_type_map.get(task.taskType)
                if alert_sender_class:
                    alert_sender = alert_sender_class(task=task, warnings=warnings)
                    alert_sender.send_alerts()

        # 更新套餐数据及加黑列表
        self.update_balance_and_add_black(task)
        # 推送数据到用户
        try:
            self.send_api_message(task)
        except:
            pass

    def update_balance_and_add_black(self, task: Task):
        """更新用户套餐数据"""
        # 如果是源站安全检测任务， 则不进行一下操作
        if task.name == "源站安全检测":
            return
        uid = task.uid
        record = Balances.objects.find_one({"uid": uid})
        if not record:
            app.logger.error(f"[0F][余额更新]查询余额Mongo记录失败: uid:{uid}")
            return
        if not record.planId:
            app.logger.error(f"[0F][余额更新]未购买套餐: uid:{uid}")
            return
        if record.planExpireTime < datetime.datetime.now():
            app.logger.error(f"[0F][余额更新]套餐已过期: uid:{uid}")
            return
        if task.taskType in {"vul", "securityEvent"}:
            # 云安全加黑
            try:
                system_add_black_list(task_ids=[task.id])
            except Exception as e:
                app.logger.error(f"加黑请求发送失败：task_id: {task.id}, error: {e}")
            else:
                app.logger.info(f"加黑请求发送成功: task_id: {task.id}")

            if self.retry:
                return

            type_max_page_key = f"{task.taskType}MaxPage"
            # 兼容statistics.urls为null的情况, 以及KeyError的情况
            try:
                inc_value = task.result.addition.get("statistics", {}).get("urls", 1) or 1
            except (KeyError, AttributeError):
                inc_value = 1
            update_result = Balances._get_collection().update_one(
                {
                    "uid": uid
                },
                {
                    "$inc": {
                        "balance.maxPage.used": inc_value,
                        f"balance.{type_max_page_key}.used": inc_value,
                    },

                }
            )
            return update_result

    def warning_counts_update(self, warnings):
        res = {
            "warningCount": len([w for w in warnings if not w.ignoredTime]),
            "warningResolvedCount": len([w for w in warnings if not w.ignoredTime and w.isResolved])
        }
        res["warningUnresolvedCount"] = res["warningCount"] - res["warningResolvedCount"]

        return res


if __name__ == '__main__':
    from service.run import app
    a = TaskResultHandler().start()
    print(a)
